#!/usr/bin/python

import sys,os
from pylab import *
from scipy.ndimage import filters,morphology,interpolation
from scipy import stats
import tables
import multiprocessing
from pdb import pm
from llpy.sutils import *

import argparse
parser = argparse.ArgumentParser(description="Find characters that are not matched well by prototypes+sigmas+dstats.")
parser.add_argument('-d','--data',default='DATA/train.h5',help="prototype file")
parser.add_argument('-p','--protos',default='DATA/protos032a.h5',help="protos file")
parser.add_argument('-o','--output',default='exceptions.h5',help="output file")
parser.add_argument('-c','--byclass',action="store_true",help="perform exception processing by class")
parser.add_argument('-R','--radius',default=None,type=float,help="absolute radius for thresholding (overrides -t)")
parser.add_argument('-t','--threshold',default=3.0,type=float,help="threshold for filtering characters out")
parser.add_argument('-N','--nsamples',default=2000000000,type=int,help="number of samples")
parser.add_argument('-Q','--par',default=multiprocessing.cpu_count(),type=int,help="parallelism")
parser.add_argument('-i','--invert',action="store_true",help="invert the selection")
parser.add_argument('-r','--pattern',default='[^~_]',help="select patterns")
#parser.add_argument('-r','--rejects',action="store_true",help="include rejects")
parser.add_argument('--chunksize',default=10000,type=int)
args = parser.parse_args()
#args = parser.parse_args(["--chunksize","1000","-Q","1","-p","tstats.h5","data/dtraining.h5"])

protos = Protos()
protos.load(args.protos)
if protos.sigmas is None:
    print "using nearest neighbor"
else:
    print "using likelihood"
if args.radius is None:
    if protos.d_means is None:
        print "d_means/d_sigmas are needed when using -t threshold"
        sys.exit(1)

def process(job):
    print job
    start,end = job
    rel = None
    with tables.openFile(args.data) as db:
        classes = db.root.classes[start:end]
        data = array(db.root.patches[start:end],'f')
        sel = array([i for i,c in enumerate(classes) if re.match('^'+args.pattern+'$',udecode(c))])
        if len(data)==0: print "warning: no characters selected; maybe use -p"; return
        data = data[sel]
        classes = classes[sel]
        indexes = find(sel)
        # if "rel" in dir(db.root): rel = db.root.rel[start:end]
        if args.byclass:
            neighbors,dists = protos.nn_index(make2d(data),100)
            n1s = []
            d1s = []
            for i,row in enumerate(neighbors):
                rcs = classes[row]
                match = find(rcs==classes[i])
                if len(match)<1:
                    n1s.append(0) # use prototype 0 to make the rest work out
                    d1s.append(999999999.0) # pick a huge distance
                else:
                    n1s.append(neighbors[match[0]])
                    d1s.append(dists[match[0]])
            n1s = array(n1s)
            d1s = array(d1s)
        else:
            neighbors,dists = protos.nn_index(make2d(data),2)
            n1s = array(neighbors)[:,0]
            d1s = array(dists)[:,0]
        if args.radius is not None:
            if args.invert:
                bad = find(d1s<=args.radius)
            else:
                bad = find(d1s>args.radius)
        else:
            means = array([protos.d_means[i] for i in neighbors],'f').ravel() # FIXME
            sigmas = array([protos.d_sigmas[i] for i in neighbors],'f').ravel() # FIXME
            c1s = (d1s-means)**2/(2.0*maximum(sigmas,1e-2)**2)
            if args.invert:
                bad = find(c1s<=args.threshold)
            else:
                bad = find(c1s>args.threshold)
    print job,"got",len(bad)
    with flock(args.output+".lock",1):
        with tables.openFile(args.output,"r+") as odb:
            for b in bad:
                odb.root.patches.append([data[b]])
                odb.root.classes.append([classes[b]])
                #if rel is not None: odb.root.rel.append(rel)

with tables.openFile(args.data) as db, tables.openFile(args.output,"w") as odb:
    nsamples = min(args.nsamples,len(db.root.classes))
    shape = list(db.root.patches.shape[1:])
    table_lcopy(db,odb)
    table_log(odb,str(sys.argv))
    print "creating output array"
    odb.createEArray(odb.root,"patches",tables.Float32Atom(),shape=[0]+shape,filters=tables.Filters(9))
    odb.createEArray(odb.root,"classes",tables.Int64Atom(),shape=[0],filters=tables.Filters(9))
    # if "rel" in dir(db.root): odb.createEArray(odb.root,"rel",tables.Float32Atom(),shape=[0,3],filters=tables.Filters(9))

jobs = chunks(nsamples,args.chunksize)
list(poolmap(process,jobs))
